
import numpy as np
import sophon.sail as sail
from easydict import EasyDict as edict
import cv2
import time
import logging
from pathlib import Path
from PIL import Image
import yaml
import argparse

from utils import set_logger
import time

from detector import Detector

def main(args):
    logging.info('main start')
    image_save_dir = Path(args.image_save_dir)
    image_save_dir.mkdir(parents=True, exist_ok=True)
    
    detector = Detector(
        args.bmodel, args.conf_thres, 0.4,
        min_box=args.get('min_box', (0, 0)),
        extra_scales=args.get('extra_scales', ())
    )
    logging.info('detectot build finished.')

    if args.input.endswith('.mp4'):
        cap = cv2.VideoCapture(args.input)
        assert cap.isOpened()
        total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
        cur_idx = cap.get(cv2.CAP_PROP_POS_FRAMES)
        frame_rate = cap.get(cv2.CAP_PROP_FPS)
        while cap.isOpened() and cur_idx < total_frames:
            cap.set(cv2.CAP_PROP_POS_FRAMES, cur_idx)
            success, frame = cap.read()
            if not success:
                logging.warn('getFrame fails')
                time.sleep(1)
                continue
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            boxes, confs = detector(rgb_frame)
            logging.info('detect: %s' % len(boxes))
            for box in boxes:
                box = box.astype(np.int)
                cv2.rectangle(rgb_frame, (box[0], box[1]), (box[2], box[3]), (255, 0, 0))
            
            out_file = image_save_dir / ('%06d.jpg' % cur_idx)
            logging.info('out_file: %s' % out_file)
            Image.fromarray(rgb_frame).save(str(out_file))
            cur_idx += frame_rate * args.sample_interval
        logging.info('camera is closed...exit')
            
    else:
        if Path(args.input).is_file():
            inputs = [Path(args.input)]
        else:
            inputs = []
            inputs.extend(list(Path(args.input).glob('**/*.jpg')))
            inputs.extend(list(Path(args.input).glob('**/*.png')))
        for f in inputs:
            frame = cv2.imread(str(f))            
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            tic = time.time()
            boxes, confs = detector(rgb_frame)
            toc = time.time()
            logging.info('detect: %s, time: %.3f secs' % (len(boxes), toc-tic))
            tic = time.time()
            for box in boxes:
                box = box.astype(np.int)
                cv2.rectangle(rgb_frame, (box[0], box[1]), (box[2], box[3]), (255, 0, 0))
            
            out_file = image_save_dir / (f.stem + '.jpg')
            logging.info('out_file: %s' % out_file)
            Image.fromarray(rgb_frame).save(str(out_file))
            toc = time.time()
            logging.info('draw and save time: %.3f secs' % (toc-tic))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('config_file', default='./demo.yaml', type=str, nargs='?', help='.yaml config file')
    cmd_args = parser.parse_args()

    args = yaml.load(open(cmd_args.config_file, 'r'), Loader=yaml.FullLoader)
    args = edict(args)

    set_logger()
    main(args)
